{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "3bdb1f74-fcdc-491a-9d8e-5de664d43ab4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import torch\n",
    "cam_root = '/data/gq/data/imagenet/nips/gradcam/samples/'\n",
    "# cam = cv2.imread(cam_root+'0aebe24fc257286e'+'.png', cv2.COLOR_BGR2GRAY) / 255.\n",
    "cam = cv2.imread(cam_root+'0aebe24fc257286e'+'.png', 0) / 255.\n",
    "cam = cv2.resize(cam, (64, 64))\n",
    "cam = torch.tensor(cam).float()\n",
    "# cam = cam.permute(2, 0, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "bd2e7c2b-1195-4a49-a207-82630d733b82",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2235, 0.2235, 0.2235,  ..., 0.0157, 0.0157, 0.0157],\n",
       "        [0.2235, 0.2235, 0.2235,  ..., 0.0157, 0.0157, 0.0157],\n",
       "        [0.2235, 0.2235, 0.2235,  ..., 0.0157, 0.0157, 0.0157],\n",
       "        ...,\n",
       "        [0.0706, 0.0706, 0.0706,  ..., 0.0667, 0.0667, 0.0667],\n",
       "        [0.0706, 0.0706, 0.0706,  ..., 0.0667, 0.0667, 0.0667],\n",
       "        [0.0706, 0.0706, 0.0706,  ..., 0.0667, 0.0667, 0.0667]])"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "17c01317-d313-49bb-a474-0abcde8fc5e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = torch.zeros((64, 64))\n",
    "(x, y) = torch.randint(0, 57, (2,))\n",
    "mask[x:x+7, y:y+7] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "21d2df70-d28a-42d8-971e-f36e80384264",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.multinomial(cam, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "08f7c6aa-9eea-4356-afcf-af646627c69e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(64, 64)\n",
      "采样到的坐标： 31 18\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "def sample_coordinates(prob_matrix):\n",
    "    # 将概率矩阵展平为一维数组\n",
    "    flattened_probs = prob_matrix.flatten()\n",
    "    \n",
    "    # 使用numpy.random.choice进行采样\n",
    "    # 参数p指定概率数组，size指定采样的数量\n",
    "    sampled_index = np.random.choice(flattened_probs.size, p=flattened_probs)\n",
    "    \n",
    "    # 将一维索引转换回二维坐标\n",
    "    sampled_coordinates = np.unravel_index(sampled_index, prob_matrix.shape)\n",
    "    \n",
    "    return sampled_coordinates\n",
    "\n",
    "# 生成一个示例的概率矩阵\n",
    "cam = torch.clamp(cam+0.2, 0.3, 0.7)\n",
    "prob_matrix = cam.numpy()\n",
    "print(prob_matrix.shape)\n",
    "prob_matrix /= prob_matrix.sum()  # 确保概率之和为1\n",
    "# 进行采样\n",
    "(x, y) = sample_coordinates(prob_matrix)\n",
    "\n",
    "print(\"采样到的坐标：\", x, y)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "5a86720c-377c-4950-8a58-b6eb1f184a84",
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = torch.zeros((64, 64))\n",
    "x_left, x_right = max(0, x-3), min(64, x+4)\n",
    "y_left, y_right = max(0, y-3), min(64, y+4)\n",
    "mask[x_left:x_right, y_left:y_right] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "0023dcab-1fbb-49dc-94fd-9f478f47453a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 1., 1., 1., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 0., 0., 0.],\n",
       "        [1., 1., 1., 1., 0., 0., 0.]])"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mask[31:35, 18:25]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93ab84b8-fe2c-43ce-8f3d-71363568a33e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python(diffusion)",
   "language": "python",
   "name": "diffusion"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
